import torch
from datasets import load_dataset, Dataset
from modelscope import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig, DataCollatorForCompletionOnlyLM

# 加载原始数据集
raw_ds = load_dataset(
    "json",
    data_files={"train": "cat.json"},
    split="train"
)

# 转换为对话格式
convs = []
for item in raw_ds:
    convs.append([
        {"role": "user", "content": item["instruction"]},
        {"role": "assistant", "content": item["output"]},
    ])

# 创建对话格式Dataset
raw_conv_ds = Dataset.from_dict({"conversations": convs})

# 加载模型和分词器
model_name = "Qwen/Qwen3-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 应用聊天模板生成训练文本
chat_inputs = tokenizer.apply_chat_template(
    raw_conv_ds["conversations"],
    tokenize=False
)

# 转换为带有"text"字段的Dataset
train_dataset = Dataset.from_dict({"text": chat_inputs})

# 创建DataCollator
response_template = "<|im_start|>assistant\n"
collator = DataCollatorForCompletionOnlyLM(
    response_template=response_template,
    tokenizer=tokenizer,
    mlm=False
)